package sshserver

import (
	"awesomeProject/goAgentServer/db"
	"awesomeProject/goAgentServer/encryption"
	"errors"
	"fmt"
	"github.com/gliderlabs/ssh"
	"github.com/pkg/sftp"
	cryptoSSH "golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"
)

type listerat []os.FileInfo

type FileLister interface {
	Filelist(*sftp.Request) (sftp.ListerAt, error)
}

type sftpSessions struct {
	session        io.ReadWriteCloser
	users          string
	sftpClient     *sftp.Client
	lock           sync.RWMutex
	sessionLog     *os.File
	backendMachine uint
}

func (f listerat) ListAt(ls []os.FileInfo, offset int64) (int, error) {
	var n int
	if offset >= int64(len(f)) {
		return 0, io.EOF
	}
	n = copy(ls, f[offset:])
	if n < len(ls) {
		return n, io.EOF
	}
	return n, nil
}

func sftpHandler(sess ssh.Session) {
	defer log.Println(sess.User(), "disconnect sftp...")
	log.Println(sess.User(), "connect sftp...")

	var u db.Users
	u = db.QueryUserInfo(sess.User())
	if u.LastHostID == 0 {
		log.Println("open sftp errors,user last host id is 0,users:", sess.User())
		return
	}

	var h db.Hosts
	h, err := db.QueryHostsInfo(u.LastHostID)
	if err != nil {
		log.Println("open sftp error,query hosts errors", err)
		return
	}

	var sshClient *cryptoSSH.Client
	if len(h.Key) == 0 {
		password, err := encryption.Decrypt(h.Password)
		if err != nil {
			log.Println("encryption.Decrypt hostInfo.Password error ,err: ", err)
			return
		}
		sshClient, err = cryptoSSH.Dial("tcp", h.Addr, &cryptoSSH.ClientConfig{
			User:            h.Hostname,
			Auth:            []cryptoSSH.AuthMethod{cryptoSSH.Password(password)},
			HostKeyCallback: func(hostname string, remote net.Addr, key cryptoSSH.PublicKey) error { return nil },
		})
	} else {
		keys, err := encryption.Decrypt(h.Key)
		if err != nil {
			log.Println("encryption.Decrypt hostInfo.Key error ,err: ", err)
			return
		}
		singer, err := cryptoSSH.ParsePrivateKey([]byte(keys))
		if err != nil {
			log.Println("ssh parse private key error ", err)
			return
		}
		sshClient, err = cryptoSSH.Dial("tcp", h.Addr, &cryptoSSH.ClientConfig{
			User:            h.Hostname,
			Auth:            []cryptoSSH.AuthMethod{cryptoSSH.PublicKeys(singer)},
			HostKeyCallback: func(hostname string, remote net.Addr, key cryptoSSH.PublicKey) error { return nil },
		})
	}
	if err != nil {
		log.Println("ssh dial error")
		return
	}

	sftpClient, err := sftp.NewClient(sshClient)
	if err != nil {
		log.Println("sftp new client error")
		return
	}

	defer sftpClient.Close()
	defer sshClient.Close()

	connection := &sftpSessions{
		session:    sess,
		users:      sess.User(),
		sftpClient: sftpClient,
	}

	server := sftp.NewRequestServer(
		sess,
		sftp.Handlers{
			FileGet:  connection,
			FilePut:  connection,
			FileCmd:  connection,
			FileList: connection,
		},
		sftp.WithStartDirectory(SftpDir))

	sessionLogName := fmt.Sprintf("%s/sftp_linux_session_%s_%s_%v.log", SessionDir, sess.User(), strings.Split(h.Addr, ":")[0], time.Now().Format("2006-01-02"))
	log.Println(sess.User(), "connect hosts:", h.Id, "ip:", h.Addr, "session log filename:", sessionLogName)
	f, err := os.OpenFile(sessionLogName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Println("can not open session log", err)
		return
	}
	defer f.Close()
	connection.sessionLog = f
	connection.backendMachine = u.LastHostID

	connection.logWrite(fmt.Sprintf("%s connect the SFTP of the machine with ID %v", connection.users, connection.backendMachine))
	defer connection.logWrite(fmt.Sprintf("%s end the sftp connection for the machine with id %v", connection.users, connection.backendMachine))

	err = server.Serve()
	if err != nil {
		if err == io.EOF {
			return
		}
		log.Println(err)
	}
}

func (c *sftpSessions) logWrite(sftpLogs string) {
	nowTime := time.Now().Format("2006-01-02 15:04:05")
	c.sessionLog.WriteString(fmt.Sprintf("%s %s\n", nowTime, sftpLogs))
}

func (c *sftpSessions) Fileread(request *sftp.Request) (io.ReaderAt, error) {

	c.logWrite(fmt.Sprintf("%s starts FileRead on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))
	defer c.logWrite(fmt.Sprintf("%s ends FileRead on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))

	c.lock.Lock()
	defer c.lock.Unlock()

	var filePath string

	filePath = request.Filepath

	file, err := c.sftpClient.Open(filePath)
	if err != nil {
		return nil, err
	}
	return file, nil
}

func (c *sftpSessions) Filewrite(request *sftp.Request) (io.WriterAt, error) {

	c.logWrite(fmt.Sprintf("%s starts Filewrite on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))
	defer c.logWrite(fmt.Sprintf("%s ends Filewrite on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))

	c.lock.Lock()
	defer c.lock.Unlock()

	var filePath string

	filePath = request.Filepath

	file, err := c.sftpClient.Create(filePath)
	if err != nil {
		return nil, err
	}
	return file, nil

}

func (c *sftpSessions) Filecmd(request *sftp.Request) error {

	c.logWrite(fmt.Sprintf("%s starts Filecmd on machine id %v, path: %s , method: %s ", c.users, c.backendMachine, request.Filepath, request.Method))
	c.logWrite(fmt.Sprintf("%s ends Filecmd on machine id %v, path: %s , method: %s ", c.users, c.backendMachine, request.Filepath, request.Method))

	c.lock.Lock()
	defer c.lock.Unlock()

	command := request.Method

	var filePath string

	filePath = request.Filepath

	switch command {
	case "Remove":
		err := c.sftpClient.Remove(filePath)
		if err != nil {
			return nil
		}
	case "Rmdir":
		err := c.sftpClient.RemoveDirectory(filePath)
		if err != nil {
			return nil
		}
	case "Mkdir":
		err := c.sftpClient.Mkdir(filePath)
		if err != nil {
			return err
		}
	default:
		return errors.New("unsupported command")
	}

	return nil
}

func (c *sftpSessions) Filelist(request *sftp.Request) (sftp.ListerAt, error) {

	c.logWrite(fmt.Sprintf("%s starts Filelist on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))
	defer c.logWrite(fmt.Sprintf("%s ends Filelist on machine id %v, path: %s", c.users, c.backendMachine, request.Filepath))

	c.lock.Lock()
	defer c.lock.Unlock()

	var dirPath string

	dirPath = request.Filepath

	dir, err := c.sftpClient.ReadDir(dirPath)
	if err != nil {
		return nil, err
	}
	return listerat(dir), err
}
